import argparse
import os.path as osp
import random
import time
import sys
import torch
from torch.nn import Linear, ReLU, ModuleList, Softmax, CrossEntropyLoss
from torch.optim.lr_scheduler import ReduceLROnPlateau
from torch_geometric.nn import GCNConv, BatchNorm, MessagePassing, global_mean_pool ,global_max_pool, LEConv
from torch_geometric.data import DataLoader
import torch.nn.functional as F
from .overloader import overload


class SPMotifNet(torch.nn.Module):
    # 这是Python类中的特殊方法 __init__()，它在创建对象时被调用。self 是一个对象本身的引用，而其他参数是初始化模型时所需的参数。
    # in_channels 是输入特征的通道数，hid_channels 是隐藏层的通道数（默认为 64），num_classes 是输出类别的数量（默认为 3），n
    # um_unit 是局部节点嵌入卷积层的数量（默认为 2）。
    def __init__(self, in_channels, hid_channels=64, num_classes=3, num_unit=2):
        super().__init__()

        # 这一行将传入的 num_unit 参数赋值给模型的 num_unit 属性，用于控制局部节点嵌入卷积层的数量。
        self.num_unit = num_unit
        # 这一行创建了一个线性层 (Linear)，用于将输入特征映射到隐藏层特征空间。它将输入特征的维度由 in_channels 映射到 hid_channels 维度。
        self.node_emb = Linear(in_channels, hid_channels)

        # 这两行分别创建了空的模块列表，用于存储局部节点嵌入卷积层 (LEConv) 和 ReLU 激活函数层。
        self.convs = ModuleList()
        self.relus = ModuleList()
        # 这一行开启一个循环，循环次数由 num_unit 决定，用于创建多个局部节点嵌入卷积层和 ReLU 激活函数层。
        for i in range(num_unit):
            # 这一行创建了一个局部节点嵌入卷积层 (LEConv)，它的输入和输出通道数均为 hid_channels。
            conv = LEConv(in_channels=hid_channels, out_channels=hid_channels)
            # 这两行将创建的局部节点嵌入卷积层和 ReLU 激活函数层添加到相应的模块列表中。
            self.convs.append(conv)
            self.relus.append(ReLU())

        # 这一行创建了一个序列模型 (Sequential)，它由几个线性层和 ReLU 激活函数组成，用于因果预测。序列模型会按顺序执行每个层。
        self.causal_mlp = torch.nn.Sequential(
            # 这一行创建了一个线性层，将隐藏层特征映射到两倍隐藏层特征维度，用于增加模型的表示能力。
            Linear(hid_channels, 2*hid_channels),
            # 这一行创建了一个 ReLU 激活函数层，用于引入非线性变换。
            ReLU(),
            # 这一行创建了一个线性层，将两倍隐藏层特征维度映射到输出类别数量 num_classes 的维度，用于输出预测结果。
            Linear(2*hid_channels, num_classes)
        )

        # 这一行创建了一个序列模型，用于置信度预测。它包含了类似于因果预测的结构。
        # 这部分的代码逻辑与因果预测的部分类似，只是输出的维度为 3（表示置信度有三种可能性）。
        self.conf_mlp = torch.nn.Sequential(
            Linear(hid_channels, 2*hid_channels),
            ReLU(),
            Linear(2*hid_channels, 3)
        )
        # 这一行创建了一个线性层，用于处理置信度预测结果。
        self.cq = Linear(3, 3)
        # 这一行创建了一个序列模型，用于处理置信度预测结果。它包含了之前定义的置信度预测部分和一个额外的线性层。
        self.conf_fw = torch.nn.Sequential(
            self.conf_mlp,
            self.cq
        )

    # 这是模型的前向传播方法。它接收四个参数：x 是输入特征张量，edge_index 是边的索引张量，edge_attr 是边的特征张量，batch 是图的分区张量。
    @overload
    def forward(self, x, edge_index, edge_attr, batch):
        # 这一行调用了 get_node_reps() 方法，用于计算节点表示。
        # 它将输入特征 x、边的索引 edge_index、边的特征 edge_attr 和图的分区 batch 作为参数传递给 get_node_reps() 方法，并将结果存储在 node_x 中。
        node_x = self.get_node_reps(x, edge_index, edge_attr, batch)
        # 这一行使用全局平均池化函数 global_mean_pool() 计算图的表示。
        # 它将节点表示 node_x 和图的分区 batch 作为参数传递给 global_mean_pool() 函数，并将结果存储在 graph_x 中。
        graph_x = global_mean_pool(node_x, batch)
        # 这一行调用了 get_causal_pred() 方法，用于预测因果性。它将图的表示 graph_x 作为参数传递给 get_causal_pred() 方法，并返回预测结果。
        return self.get_causal_pred(graph_x)

    # 这是 get_node_reps() 方法的定义，它接收四个参数：x 是输入特征张量，edge_index 是边的索引张量，edge_attr 是边的特征张量，batch 是图的分区张量。
    @overload
    def get_node_reps(self, x, edge_index, edge_attr, batch):
        # 这一行将输入特征 x 通过节点嵌入层 node_emb 进行线性变换。node_emb 是在 __init__() 方法中初始化的一个线性层，用于将输入特征映射到隐藏层特征空间。
        x = self.node_emb(x)
        # 这一行使用 Python 的 zip() 函数将 convs 和 relus 两个列表中的元素一一配对，然后进行迭代。
        # convs 列表包含了多个局部节点嵌入卷积层，relus 列表包含了相应数量的 ReLU 激活函数层。
        for conv, ReLU in \
                zip(self.convs, self.relus):
            # 这一行调用了当前迭代的局部节点嵌入卷积层 conv，将输入特征 x、边的索引 edge_index 和边的特征 edge_attr 作为参数传递给该卷积层，并将结果存储在 x 中。
            x = conv(x=x, edge_index=edge_index, edge_weight=edge_attr)
            # 这一行将 ReLU 激活函数应用到 x 上，对每个元素进行非线性变换。
            x = ReLU(x)
        # 这一行将经过多次卷积和 ReLU 操作后的节点特征表示存储在 node_x 变量中。
        node_x = x
        # 这一行返回经过局部节点嵌入卷积层和 ReLU 激活函数层处理后的节点特征表示。
        return node_x

    # get_graph_rep() 方法实现了从节点表示计算到图表示的过程，通过全局平均池化操作得到最终的图的表示。
    # 这是 get_graph_rep() 方法的定义，它接收四个参数：x 是输入特征张量，edge_index 是边的索引张量，edge_attr 是边的特征张量，batch 是图的分区张量。
    @overload
    def get_graph_rep(self, x, edge_index, edge_attr, batch):
        # 这一行调用了 get_node_reps() 方法，用于计算节点表示。
        # 它将输入特征 x、边的索引 edge_index、边的特征 edge_attr 和图的分区 batch 作为参数传递给 get_node_reps() 方法，并将结果存储在 node_x 中。
        node_x = self.get_node_reps(x, edge_index, edge_attr, batch)
        # 这一行使用全局平均池化函数 global_mean_pool() 计算图的表示。
        # 它将节点表示 node_x 和图的分区 batch 作为参数传递给 global_mean_pool() 函数，并将结果存储在 graph_x 中。
        graph_x = global_mean_pool(node_x, batch)
        # 这一行返回计算得到的图的表示 graph_x。
        return graph_x

    # get_causal_pred() 方法实现了对因果图的表示进行预测的过程，通过多层感知器模型将因果图的表示映射到预测结果空间。
    # 这是 get_causal_pred() 方法的定义，它接收两个参数：self 是类的实例，causal_graph_x 是因果图的表示。
    def get_causal_pred(self, causal_graph_x):
        # 这一行调用了因果预测的多层感知器模型 (causal_mlp)，对因果图的表示 causal_graph_x 进行前向传播，得到预测结果 pred。
        # 在初始化模型时，causal_mlp 被定义为一个序列模型，其中包含了线性层和 ReLU 激活函数，用于将因果图的表示映射到最终的预测结果空间。
        pred = self.causal_mlp(causal_graph_x)
        # 这一行返回预测结果 pred，即因果性预测的输出。
        return pred

    # get_conf_pred() 方法实现了对置信度图的表示进行预测的过程，通过序列模型将置信度图的表示映射到预测结果空间。
    # 这是 get_conf_pred() 方法的定义，它接收两个参数：self 是类的实例，conf_graph_x 是置信度图的表示。
    def get_conf_pred(self, conf_graph_x):
        # 这一行调用了置信度预测的前向模型 (conf_fw)，将置信度图的表示 conf_graph_x 作为输入，进行前向传播，得到预测结果 pred。
        # 在初始化模型时，conf_fw 被定义为一个序列模型，其中包含了置信度预测部分 (conf_mlp) 和一个线性层 (cq)，用于将置信度图的表示映射到最终的预测结果空间。
        pred = self.conf_fw(conf_graph_x)
        # 这一行返回预测结果 pred，即置信度预测的输出。
        return pred

    # get_comb_pred() 方法实现了对因果图和置信度图的表示进行组合预测的过程，将置信度预测结果加权应用到因果性预测结果上，得到最终的组合预测结果。
    # 这是 get_comb_pred() 方法的定义，它接收三个参数：self 是类的实例，causal_graph_x 是因果图的表示，conf_graph_x 是置信度图的表示。
    def get_comb_pred(self, causal_graph_x, conf_graph_x):
        # 这一行调用了因果预测的多层感知器模型 (causal_mlp)，对因果图的表示 causal_graph_x 进行前向传播，得到因果性预测结果 causal_pred。
        # 这个结果表示模型对于给定因果图的预测结果。
        causal_pred = self.causal_mlp(causal_graph_x)
        # 这一行调用了置信度预测的多层感知器模型 (conf_mlp)，对置信度图的表示 conf_graph_x 进行前向传播，得到置信度预测结果 conf_pred。
        # .detach() 方法用于断开 conf_pred 与计算图的连接，将其从计算图中分离出来，以防止对置信度预测结果的梯度传播影响因果预测结果。
        conf_pred = self.conf_mlp(conf_graph_x).detach()
        # 这一行将置信度预测结果 conf_pred 经过 sigmoid 函数进行归一化（将其值缩放到 [0, 1] 区间），然后与因果性预测结果 causal_pred 相乘。
        # 这样做的目的是使用置信度预测结果对因果性预测结果进行加权，从而得到最终的组合预测结果。
        # 这一行返回组合预测结果，即使用置信度预测结果对因果性预测结果进行加权后得到的最终预测结果。
        return torch.sigmoid(conf_pred) * causal_pred

    # 这是 reset_parameters() 方法的定义，它用于重新初始化模型的参数。
    def reset_parameters(self):
        # 这一行使用了 PyTorch 中的 torch.no_grad() 上下文管理器，它可以禁用梯度的计算。这意味着在接下来的代码块中，不会对参数的梯度进行追踪和更新，只会执行参数值的修改操作。
        with torch.no_grad():
            # 这一行迭代模型的所有参数。self.parameters() 方法返回模型中的所有参数。
            for param in self.parameters():
                # 这一行对参数进行重新初始化。调用了参数的 uniform_() 方法，该方法将参数的值从均匀分布中采样，并将其值设置为介于 -1.0 和 1.0 之间。
                # 这样做的目的是在训练开始之前，将模型的参数随机初始化到一个合适的范围内，以提高模型的训练稳定性和收敛速度。
                param.uniform_(-1.0, 1.0)
